#!/usr/bin/env python
# -*- coding: UTF-8 -*-
# @Project:     data_manager
# @File:        models.py
# @Author:      Fan GuiChuan
# @Date:        2024/12/17 9:15
# @Version:     Python3.7
# @Description:
# my_db_tool/models.py
import inspect

from peewee import SqliteDatabase, MySQLDatabase, PostgresqlDatabase
from playhouse.reflection import Introspector


def save_models_to_file(models, file_path):
    """将生成的模型保存为 Python 文件"""
    with open(file_path, 'w') as f:
        f.write("from peewee import *\n\n")
        for model_name, model_class in models.items():
            f.write(f"class {model_name}(Model):\n")
            for line in inspect.getsourcelines(model_class)[0][1:]:
                f.write(f"    {line}")
            f.write("\n\n")
    print(f"模型已保存到文件：{file_path}")


def get_database(db_config):
    """根据数据库类型获取数据库连接"""
    db_type = db_config.get('type')

    if db_type == 'sqlite':
        return SqliteDatabase(db_config['name'])
    elif db_type == 'mysql':
        return MySQLDatabase(
            db_config['name'],
            user=db_config['user'],
            password=db_config['password'],
            host=db_config['host'],
            port=db_config['port']
        )
    elif db_type == 'postgres':
        return PostgresqlDatabase(
            db_config['name'],
            user=db_config['user'],
            password=db_config['password'],
            host=db_config['host'],
            port=db_config['port']
        )
    else:
        raise ValueError(f"不支持的数据库类型：{db_type}")


def generate_models_from_database(db):
    """从数据库生成模型并保存到文件"""
    introspector = Introspector.from_database(db)
    models = introspector.generate_models()
    return models


def get_tables(db):
    """获取数据库中所有表的名称"""
    introspector = Introspector.from_database(db)
    models = introspector.generate_models()
    return list(models.keys())